//! 冷池管理
//! 
//! 存储Firecracker快照模板，启动时间~200ms

use crate::Result;
use crate::types::{PoolStats, PoolType, Language, SandboxSpec};
use crate::sandbox_manager::{Sandbox, LifecycleManager};
use crate::firecracker::snapshot::SnapshotManager;
use std::sync::Arc;
use std::collections::HashMap;
use tokio::sync::{RwLock, Mutex};
use std::time::Instant;
use tracing::{info, warn, error};

/// 快照信息
#[derive(Debug, Clone)]
pub struct SnapshotInfo {
    pub snapshot_id: String,
    pub language: Language,
    pub path: String,
    pub memory_file_path: String,
    pub kernel_image_path: String,
    pub rootfs_path: String,
    pub created_at: std::time::SystemTime,
    pub size_mb: u64,
}

/// 冷池管理器
pub struct ColdPool {
    config: Arc<crate::config::PoolConfig>,
    /// 快照存储映射
    snapshots: Arc<RwLock<HashMap<Language, Vec<SnapshotInfo>>>>,
    /// 快照管理器
    snapshot_manager: Arc<Mutex<SnapshotManager>>,
    /// 生命周期管理器
    lifecycle_manager: Arc<LifecycleManager>,
    /// 正在创建的沙箱计数
    creating_count: Arc<std::sync::atomic::AtomicUsize>,
}

impl ColdPool {
    /// 创建冷池管理器
    pub fn new(
        config: Arc<crate::config::PoolConfig>,
        snapshot_manager: Arc<Mutex<SnapshotManager>>,
        lifecycle_manager: Arc<LifecycleManager>,
    ) -> Self {
        Self {
            config,
            snapshots: Arc::new(RwLock::new(HashMap::new())),
            snapshot_manager,
            lifecycle_manager,
            creating_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        }
    }
    
    /// 从快照创建沙箱
    pub async fn acquire(&self, language: &Language) -> Result<Option<Sandbox>> {
        let start_time = Instant::now();
        info!("开始从冷池快照创建沙箱 - 语言: {:?}", language);
        
        // 查找对应语言的快照
        let snapshot_info = {
            let snapshots = self.snapshots.read().await;
            snapshots.get(language)
                .and_then(|snapshots| snapshots.first())
                .cloned()
        };
        
        let snapshot_info = match snapshot_info {
            Some(info) => info,
            None => {
                warn!("未找到语言 {:?} 的快照模板", language);
                return Ok(None);
            }
        };
        
        // 增加创建计数
        self.creating_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        
        // 从快照恢复VM并创建沙箱
        let create_result = self.create_sandbox_from_snapshot(snapshot_info, language).await;
        
        // 减少创建计数
        self.creating_count.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
        
        match create_result {
            Ok(sandbox) => {
                let duration = start_time.elapsed();
                info!(
                    "从冷池创建沙箱成功 - 语言: {:?}, 耗时: {}ms",
                    language,
                    duration.as_millis()
                );
                Ok(Some(sandbox))
            }
            Err(e) => {
                error!("从冷池创建沙箱失败 - 语言: {:?}, 错误: {}", language, e);
                Err(e)
            }
        }
    }
    
    /// 初始化语言快照模板
    pub async fn initialize_snapshots(&self) -> Result<()> {
        info!("初始化冷池快照模板");
        
        let languages = vec![
            Language::Python,
            Language::JavaScript,
            Language::Rust,
            Language::Go,
            Language::Java,
        ];
        
        let mut snapshots = self.snapshots.write().await;
        
        for language in languages {
            // 为每种语言创建基础快照
            let snapshot_info = self.create_base_snapshot(&language).await?;
            snapshots.entry(language)
                .or_insert_with(Vec::new)
                .push(snapshot_info);
        }
        
        info!("冷池快照模板初始化完成");
        Ok(())
    }
    
    /// 获取冷池统计
    pub async fn get_stats(&self) -> Result<PoolStats> {
        let snapshots = self.snapshots.read().await;
        let total_snapshots: usize = snapshots.values().map(|v| v.len()).sum();
        let creating_count = self.creating_count.load(std::sync::atomic::Ordering::SeqCst);
        
        Ok(PoolStats {
            pool_type: PoolType::Cold,
            total_capacity: total_snapshots,
            used_count: 0, // 冷池快照不被"使用"，只是被复制
            idle_count: total_snapshots,
            creating_count,
            error_count: 0, // TODO: 实现错误统计
            avg_allocation_latency_ms: 200.0, // 冷池目标延迟
            cache_hit_rate: if total_snapshots > 0 { 100.0 } else { 0.0 },
        })
    }
    
    /// 获取支持的语言列表
    pub async fn supported_languages(&self) -> Vec<Language> {
        let snapshots = self.snapshots.read().await;
        snapshots.keys().cloned().collect()
    }
    
    /// 获取特定语言的快照数量
    pub async fn snapshot_count(&self, language: &Language) -> usize {
        let snapshots = self.snapshots.read().await;
        snapshots.get(language).map(|v| v.len()).unwrap_or(0)
    }
    
    /// 清理过期快照
    pub async fn cleanup_expired_snapshots(&self) -> Result<()> {
        let mut snapshots = self.snapshots.write().await;
        let now = std::time::SystemTime::now();
        let max_age = std::time::Duration::from_secs(24 * 3600); // 24小时
        
        for (language, snapshot_list) in snapshots.iter_mut() {
            let initial_count = snapshot_list.len();
            snapshot_list.retain(|snapshot| {
                now.duration_since(snapshot.created_at)
                    .map(|age| age < max_age)
                    .unwrap_or(false)
            });
            
            let removed_count = initial_count - snapshot_list.len();
            if removed_count > 0 {
                info!("清理过期快照 - 语言: {:?}, 清理数量: {}", language, removed_count);
            }
        }
        
        Ok(())
    }
    
    // 私有方法
    
    /// 从快照创建沙箱
    async fn create_sandbox_from_snapshot(
        &self,
        snapshot_info: SnapshotInfo,
        language: &Language,
    ) -> Result<Sandbox> {
        // 创建沙箱规格
        let spec = SandboxSpec {
            language: language.clone(),
            memory_mb: self.config.default_memory_mb as u64,
            cpu_cores: self.config.default_cpu_cores,
            disk_mb: self.config.default_disk_mb as u64,
            timeout_secs: self.config.default_timeout_seconds as u64,
            network_enabled: true,
            env_vars: std::collections::HashMap::new(),
            snapshot_path: Some(snapshot_info.path.clone()),
            memory_file_path: Some(snapshot_info.memory_file_path.clone()),
            kernel_image_path: Some(snapshot_info.kernel_image_path.clone()),
            rootfs_path: Some(snapshot_info.rootfs_path.clone()),
        };
        
        // 通过生命周期管理器创建沙箱
        let _metadata = self.lifecycle_manager.create_sandbox(spec.clone()).await?;
        
        Ok(Sandbox::new(spec))
    }
    
    /// 创建基础快照
    async fn create_base_snapshot(&self, language: &Language) -> Result<SnapshotInfo> {
        info!("创建语言 {:?} 的基础快照", language);
        
        let snapshot_id = uuid::Uuid::new_v4();
        
        // 根据语言确定镜像路径
        let (kernel_path, rootfs_path) = self.get_language_images(language);
        
        // 创建快照
        let snapshot_metadata = self.snapshot_manager.lock().await.create_snapshot(
            snapshot_id,
            language.clone(),
            crate::firecracker::snapshot::SnapshotType::Language,
        ).await?;
        
        let memory_file_path = format!("{}.mem", snapshot_metadata.path);
        
        Ok(SnapshotInfo {
            snapshot_id: snapshot_metadata.id.to_string(),
            language: language.clone(),
            path: snapshot_metadata.path.clone(),
            memory_file_path,
            kernel_image_path: kernel_path,
            rootfs_path,
            created_at: std::time::SystemTime::now(),
            size_mb: 128, // TODO: 计算实际大小
        })
    }
    
    /// 获取语言对应的镜像路径
    fn get_language_images(&self, language: &Language) -> (String, String) {
        let base_path = "/opt/koda/images";
        
        match language {
            Language::Python => (
                format!("{}/vmlinux-python", base_path),
                format!("{}/rootfs-python.ext4", base_path),
            ),
            Language::JavaScript => (
                format!("{}/vmlinux-node", base_path),
                format!("{}/rootfs-node.ext4", base_path),
            ),
            Language::Rust => (
                format!("{}/vmlinux-rust", base_path),
                format!("{}/rootfs-rust.ext4", base_path),
            ),
            Language::Go => (
                format!("{}/vmlinux-go", base_path),
                format!("{}/rootfs-go.ext4", base_path),
            ),
            Language::Java => (
                format!("{}/vmlinux-java", base_path),
                format!("{}/rootfs-java.ext4", base_path),
            ),
            Language::Cpp => (
                format!("{}/vmlinux-cpp", base_path),
                format!("{}/rootfs-cpp.ext4", base_path),
            ),
            Language::Unknown => (
                format!("{}/vmlinux-generic", base_path),
                format!("{}/rootfs-generic.ext4", base_path),
            ),
        }
    }
}